/*
 * FILE: rdt_sender.cc
 * DESCRIPTION: Reliable data transfer sender.
 * NOTE: This implementation assumes there is no packet loss, corruption, or 
 *       reordering.  You will need to enhance it to deal with all these 
 *       situations.  In this implementation, the packet format is laid out as 
 *       the following:
 *       
 *       |<-  1 byte  ->|<-             the rest            ->|
 *       | payload size |<-             payload             ->|
 *
 *       The first byte of each packet indicates the size of the payload
 *       (excluding this single-byte header)
 */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <map>
#include <set>
#include <iostream>

#include "rdt_struct.h"
#include "rdt_sender.h"

std::map<int, std::string> sender_saved;
std::set<int> active_timers;
int sent_index = 0;

void start_timer(int index);

void kill_timer(int index);

void send_by_index(int index);

void receive_acknowledge(packet *pkt);

/* sender initialization, called once at the very beginning */
void Sender_Init() {
    fprintf(stdout, "At %.2fs: sender initializing ...\n", GetSimulationTime());
}

/* sender finalization, called once at the very end.
   you may find that you don't need it, in which case you can leave it blank.
   in certain cases, you might want to take this opportunity to release some 
   memory you allocated in Sender_init(). */
void Sender_Final() {
    fprintf(stdout, "At %.2fs: sender finalizing ...\n", GetSimulationTime());
}

/* event handler, called when a message is passed from the upper layer at the 
   sender */
void Sender_FromUpperLayer(struct message *msg) {
#if USE_ORG
    /* 1-byte header indicating the size of the payload */
    int header_size = 1;

    /* maximum payload size */
    int maxpayload_size = RDT_PKTSIZE - header_size;

    /* split the message if it is too big */

    /* reuse the same packet data structure */
    packet pkt;

    /* the cursor always points to the first unsent byte in the message */
    int cursor = 0;

    while (msg->size - cursor > maxpayload_size) {
        /* fill in the packet */
        pkt.data[0] = maxpayload_size;
        memcpy(pkt.data + header_size, msg->data + cursor, maxpayload_size);

        /* send it out through the lower layer */
        Sender_ToLowerLayer(&pkt);

        /* move the cursor */
        cursor += maxpayload_size;
    }

    /* send out the last packet */
    if (msg->size > cursor) {
        /* fill in the packet */
        pkt.data[0] = msg->size - cursor;
        memcpy(pkt.data + header_size, msg->data + cursor, pkt.data[0]);

        /* send it out through the lower layer */
        Sender_ToLowerLayer(&pkt);
    }
#else
    std::string str(msg->data, msg->size);

    int max_payload_size = RDT_PKTSIZE - HEADER_SIZE;
    int data_size = max_payload_size / 2 - INDEX_SIZE;

    /* Packet Structure
     *
     * 0~1 payload size
     * 2~9 index crc code
     * 10~127 data crc code
     *
     */

    while (!str.empty()) {
        std::string data = str.substr(0, data_size);
        str.erase(0, data_size);

        sent_index++; // Start from 1
        sender_saved.insert(std::make_pair(sent_index, data)); // Add to buffer

        send_by_index(sent_index);
    }

#endif
}

/* event handler, called when a packet is passed from the lower layer at the 
   sender */
void Sender_FromLowerLayer(struct packet *pkt) {
    receive_acknowledge(pkt);
}

/* event handler, called when the timer expires */
void Sender_Timeout() {
    // if (!active_timers.empty()) {
    //     std::cout << "Time up, retransmit." << active_timers.size() << std::endl;
    // }

    auto timers = active_timers;
    active_timers.clear();

    for (int timer_id : timers) {
        send_by_index(timer_id);
    }
}

void start_timer(int index) {
    if (!Sender_isTimerSet()) {
        Sender_StartTimer(TIMER_INTERVAL);
    }
    active_timers.insert(index);
}

void kill_timer(int index) {
    active_timers.erase(index);
    if (active_timers.empty() && Sender_isTimerSet()) {
        Sender_StopTimer();
    }
}

void send_by_index(int index) {
    auto it = sender_saved.find(index);
    if (it == sender_saved.end()) {
        return;
    }

    auto data = it->second;

    int max_payload_size = RDT_PKTSIZE - HEADER_SIZE;
    int data_size = max_payload_size / 2 - INDEX_SIZE;

    int real_size = data.size();
    if (real_size < data_size) {
        data.resize(data_size, 0);
    }
    int payload_size = (real_size + INDEX_SIZE) * 2;

    std::string encode_res = EncodeMessage(data, index);

    // Copy data to packet
    packet pkt{};
    memset(pkt.data, 0, RDT_PKTSIZE);

    pkt.data[0] = char(payload_size);
    memcpy(pkt.data + HEADER_SIZE, encode_res.data(), payload_size);

    // std::cout << "S " << real_size << " " << data << std::endl;

    Sender_ToLowerLayer(&pkt);

    // Add to timers
    start_timer(index);
}

void receive_acknowledge(packet *pkt) {
    int payload_size = pkt->data[0];
    int real_size = payload_size / 2 - INDEX_SIZE;

    const int max_payload_size = RDT_PKTSIZE - HEADER_SIZE;
    if (payload_size <= 0 || payload_size > max_payload_size) {
        return;
    }

    std::string str(pkt->data + HEADER_SIZE, payload_size);
    int index;

    std::string str_msg = DecodeMessage(str, index);
    if (str_msg != MESSAGE_ACK) {
        return;
    }

    // Remove from timers
    kill_timer(index);

    // std::cout << "Sender receive ack " << index << std::endl;
}

// CRC 8 Check
const uint16_t crc_key = 0x31;

uint16_t crc8_merge(unsigned char value) {
    unsigned char i, crc;

    crc = value;
    for (i = 8; i > 0; --i) {
        /* highest bit is 1 */
        if (crc & 0x80) {
            /* 0x31(x8+x5+x4+1，100110001) */
            crc = (crc << 1) ^ crc_key;
        } else {
            crc = (crc << 1);
        }
    }

    uint16_t res = uint16_t(value) << 8;
    res += crc;

    return res;
}

bool crc8_extract(uint16_t code, unsigned char *value = nullptr) {
    unsigned char i;
    uint16_t rem = code;

    for (i = 8; i > 0; --i) {
        /* highest bit is 1 */
        if (rem & (0x80 << 8)) {
            /* 0x31(x8+x5+x4+1，100110001) */
            rem = (rem << 1) ^ (crc_key << 8);
        } else {
            rem = (rem << 1);
        }
    }

    uint16_t res = code >> 8;
    if (value) {
        *value = res;
    }

    return rem == 0;
}

// Implementation of encoder and decoder

std::string EncodeMessage(const std::string &msg, int index) {
    std::string res;

    // Store index as char[] first
    char *p = (char *) (&index);
    for (int i = 0; i < INDEX_SIZE; ++i) {
        char ch = *p;
        uint16_t crc = crc8_merge(ch);

        std::string crc_str((char *) (&crc), sizeof(uint16_t));
        res.append(crc_str);

        p++;
    }

    for (char ch : msg) {
        uint16_t crc = crc8_merge(ch);

        std::string crc_str((char *) (&crc), sizeof(uint16_t));
        res.append(crc_str);
    }

    return res;
}

std::string DecodeMessage(const std::string &str, int &index) {
    if (str.size() <= INDEX_SIZE * 2) {
        return "";
    }

    std::string str_index;
    for (int i = 0; i < INDEX_SIZE; ++i) {
        std::string frag = str.substr(i * 2, 2);

        unsigned char ch;
        uint16_t code = *((uint16_t *) frag.data());

        if (!crc8_extract(code, &ch)) {
            return "";
        }

        str_index.push_back(char(ch));
    }
    index = *((int *) (str_index.data()));

    std::string data;
    int sz = str.size() / 2;
    for (int i = INDEX_SIZE; i < sz; ++i) {
        std::string frag = str.substr(i * 2, 2);

        unsigned char ch;
        uint16_t code = *((uint16_t *) frag.data());

        if (!crc8_extract(code, &ch)) {
            return "";
        }

        data.push_back(char(ch));
    }

    return data;
}